import json
from py2neo import Graph, Node, Relationship
from concurrent.futures import ThreadPoolExecutor
import threading

# 连接到Neo4j数据库
# graph = Graph("http://localhost:7474/",auth=("neo4j", "123456"))
graph = Graph("http://192.168.86.128:7474/", user="neo4j", password="12345678", name="neo4j")

# 定义处理单行数据的函数
def process_data(data):

    # 提取src_ip、dst_ip和dst_port字段
    src_ip = data["src_ip"]
    dst_ip = data["dst_ip"]
    dst_port = data["dst_port"]

    # 检查关系是否存在
    query = """
    MATCH (start:IP {ip: $src_ip})-[rel:CONNECTS_TO]->(end:IP {ip: $dst_ip})
    RETURN rel
    """

    with data_lock:
        result = graph.evaluate(query, src_ip=src_ip, dst_ip=dst_ip)
        if result is None:
            # 创建或获取src_ip节点
            src_node = Node("IP", ip=src_ip)
            graph.merge(src_node, "IP", "ip")

            # 创建或获取dst_ip节点
            dst_node = Node("IP", ip=dst_ip)
            graph.merge(dst_node, "IP", "ip")

            # 创建关系及属性
            rel = Relationship(src_node, "CONNECTS_TO", dst_node, times=1, dst_port=dst_port)
            graph.create(rel)
        else:
            # 更新关系的times属性值
            query = """
            MATCH (start:IP {ip: $src_ip})-[rel:CONNECTS_TO]->(end:IP {ip: $dst_ip})
            SET rel.times = rel.times + 1
            """
            graph.run(query, src_ip=src_ip, dst_ip=dst_ip)

# 读取data.json文件
with open("C:/Users/yinwe/Desktop/data.json", "r") as file:
    lines = file.readlines()

# 加入线程锁
data_lock = threading.Lock()

# 使用线程池处理数据
with ThreadPoolExecutor() as executor:
    for line in lines:
        data = json.loads(line)
        executor.submit(process_data, data)